"""
Robot Framework library for connecting to Mongo DB

Uses PyMongo Driver
See http://api.mongodb.org/python for more info.
Pre-installation: easy_install pymongo

Accedo TV

Xuan Xu
"""

ROBOT_LIBRARY_VERSION = 0.1

import pymongo
import json
import ConfigParser
from robot.utils import asserts
from pymongo import Connection

class MongodbLibrary:

    conn = None
    db = None
    
    def ConnectToMongoDB(self, host_url, db_name, port_number=27017):
        """ Connect to Mongo database
            Loads pymongo and connects to the MongoDB host.   
            Input:
                host_url      - Url to db host
		db_name       - database name
                port_number   - Port of db (default is 27017)

        """

	self.conn = Connection(host=host_url, port=port_number)
	print 'Connecting to host:"' + host_url +'"'
	mongodb = db_name

	print 'use db:"' + mongodb +'"'
	self.db = self.conn[mongodb]


    def CloseMongoConnection(self):
        """ Commits changes and closes connection to db. 
            Assumes that ConnectToMongoDB has been run first.
        """

        if self.conn is not None:
            self.conn.disconnect()
            self.conn.close()
        else:
            print '*WARN* No open connection to close'

    def GetMongoDBCollections(self):
        """
        Returns a list of all of the collections for the database you
        passed in on the connected MongoDB server.

        Usage is:
        | @{allCollections} | Get MongoDB Collections | DBName |
        | Log Many | @{allCollections} |
        | Should Contain | ${allCollections} | CollName |
        """
        allCollections = None

        if self.conn is not None: 
            allCollections = self.db.collection_names()
            print "| @{allCollections} | Get MongoDB Collections "
            
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

	return allCollections

    def DropMongodbCollection(self, dbCollName):
        """
        Deletes the named collection passed in from the database named.
        If the collection does not exist, no errors are thrown.

        Usage is:
        | Drop MongoDB Collection | CollectionName |
        | @{allCollections} | Get MongoDB Collections | myDB |
        | Should Not Contain | ${allCollections} | CollectionName |
        """
        
        if self.conn is not None:
            
            self.db.drop_collection('%s' % (dbCollName))
            print "| Drop MongoDB Collection | %s |" % (dbCollName)
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

    def ValidateMongodbCollection(self, dbCollName):
        """
        Returns a string of validation info. Raises CollectionInvalid if 
        validation fails.

        Usage is:
        | ${allResults} | Validate MongoDB Collection | CollectionName |
        | Log | ${allResults} |
        """
        if self.conn is not None:
            dbCollName = str(dbCollName)
            
            allResults = self.db.validate_collection('%s' % dbCollName)
            print "| ${allResults} | Validate MongoDB Collection | %s |" % (dbCollName)
            return allResults
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

    def singleDocument(self, dbCollName):
        """ Executes a SQL query and returns the first object in the resultset (if any)
            Assumes that connection to db has been created with connectToDb.
            N.B. Currently also only returns one string, i.e. you can only specify one column to be returned. 
	    Input:
		dbCollName 
	"""

        result = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName +'"'
	    table = self.db[dbCollName]
	    resp = table.find_one()    
	    print resp
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result

    def singleValue(self, dbCollName, key):
        """ Executes a SQL query and returns the first object in the resultset (if any)
            Assumes that connection to db has been created with connectToDb.
            N.B. Currently also only returns one string, i.e. you can only specify one column to be returned. 
	    Input:
		dbCollName 
		key
	"""

        result = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName +'"'
	    table = self.db[dbCollName]
	    resp = table.find_one()
	    print resp
		
		
	    #recordJSON = str(resp)
	    #print recordJSON
	    #print json.load(resp)
	    #parsed_data = json.loads(recordJSON)

	    result = resp[key]
	    
	    print result

        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result
   
    

    def singleQuery(self, dbCollName, key, value):
        """ Executes a single query and returns the first object in the resultset (if any)
            Assumes that connection to db has been created with connectToMongoDB.
            N.B. Currently also only returns one string, i.e. you can specify one line of result to be returned. 
	    Input:
		dbCollName       
		key           - search key
		value         - the value of key
        """
        result = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and find "' + key + '=' + value + '"'
	    table = self.db[dbCollName]

	    message = '{"%s": %s}' % (key, value)
	    recordJSON = dict(json.loads(message))

	    resp = table.find_one(recordJSON)
	    #resp = table.find_one({"accountId": 1})
	    print resp
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result

    def QueryAllRecords(self, dbCollName):
        """ Executes a multi query and returns all objects in the resultset (if any)
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
        """
        result = None
	record = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '"'
	    table = self.db[dbCollName]
	   
	    cnt = table.find().count()
	    cntmsg = 'Totally %d records' % (cnt)
	    print cntmsg
	    
	    for record in table.find():
	      record
	      result = "%s%s\n" % (result, record)
	    
            print result
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result

    def QueryAllRecordsCount(self, dbCollName):
        """ Executes a multi query and returns all objects in the resultset (if any)
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
        """
        result = None
	record = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '"'
	    table = self.db[dbCollName]
	   
	    cnt = table.find().count()
	    cntmsg = 'Totally %d records' % (cnt)
	    print cntmsg
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return cnt


    def QuerySomeRecords(self, dbCollName, key, value):
        """ Executes a multi query and returns all objects in the resultset (if any)
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
		key           - search key
		value         - the value of key
        """
        result = None
	record = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and find "' + key + '=' + value + '"'
	    table = self.db[dbCollName]

	    message = '{"%s": "%s"}' % (key, value)
	    
	    recordJSON = dict(json.loads(message))

	    #cnt = table.find({"hash": "68706a40f7f5c9e80afa6cc94351acfd52ede4d6"}).count()
	    cnt = table.find(recordJSON).count()
	    cntmsg = 'Totally %d records' % (cnt)
	    print cntmsg
	    
	    print message
	    for record in table.find(recordJSON):
	      record
	      result = "%s%s\n" % (result, record)
	    
            print result
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result



    def QuerySomeRecordsCount(self, dbCollName, key, value):
        """ Executes a multi query for the count and returns value in the resultset
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
		key           - search key
		value         - the value of key
        """
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and find "' + key + '=' + value + '"'
	    table = self.db[dbCollName]

	    message = '{"%s": "%s"}' % (key, value)

	    recordJSON = dict(json.loads(message))

	    cnt = table.find(recordJSON).count()
	    cntmsg = 'Totally %d records' % (cnt)
	    print cntmsg
	    
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return cnt

    def RemoveSomeRecords(self, dbCollName, key, value):
        """ Executes a multi query and returns all objects in the resultset (if any)
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
		key           - search key
		value         - the value of key
	    Output:
	        result        - How many records have been deleted
        """
        result = 0
	record = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and find "' + key + '=' + value + '"'
	    table = self.db[dbCollName]

	    message = '{"%s": "%s"}' % (key, value)
	    
	    recordJSON = dict(json.loads(message))

	    cnt = table.find(recordJSON).count()
	    cntmsg = 'Totally %d records' % (cnt)
	    print cntmsg

	    if cnt == 0:
		result = 0
	    else:
		coll = self.db['%s' % (dbCollName)]
		coll.remove(recordJSON)
		result = cnt

	    print result
	    return result
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')


    def AdvanceQuery(self, dbCollName, recordJSON):
        """ Executes a advanced query and returns all objects in the resultset (if any), user can specify the query but it must be correct syntax
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
		recordJSON           - advanced query (JSON format)
			e.x. { "field" : { $gt: value } } or { "field" : { $gt: value1, $lt: value2 } }
			for more queries, refer to http://www.mongodb.org/display/DOCS/Advanced+Queries	
        """
        result = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and execute "' + recordJSON + '"'
	    table = self.db[dbCollName]

	    recordJSON = str(recordJSON)

	    try:
		recordJSON = dict(json.loads(recordJSON))
	    except:
	        print "It's not a valid JSON"
		return result
	    else:
		print "It's a valid JSON"
		cnt = table.find(recordJSON).count()
	        cntmsg = 'Totally %d records' % (cnt)
	        print cntmsg
	    
	        #print message
	        for record in table.find(recordJSON):
	          record
	          result = "%s%s\n" % (result, record)
	    
	        print result 
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result

    def AdvanceSingleQueryByKey(self, dbCollName, recordJSON, key):
        """ Executes a advanced query and returns only first object by key in the resultset (if any), user can specify the query but it must be correct syntax
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName       
		recordJSON           - advanced query (JSON format)
			e.x. { "field" : { $gt: value } } or { "field" : { $gt: value1, $lt: value2 } }
			for more queries, refer to http://www.mongodb.org/display/DOCS/Advanced+Queries	
		key                  - search key
        """
        result = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and execute "' + recordJSON + '"'
	    table = self.db[dbCollName]

	    recordJSON = str(recordJSON)

	    try:
		recordJSON = dict(json.loads(recordJSON))
	    except:
	        print "It's not a valid JSON"
		return result
	    else:
		print "It's a valid JSON"
		cnt = table.find(recordJSON).count()
	        cntmsg = 'Totally %d records' % (cnt)
	        print cntmsg
	    
	        #print message
                resp = table.find_one(recordJSON)

	        result = resp[key]
	    
	        print result 
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

        return result


    def AdvanceRemove(self, dbCollName, recordJSON):
        """ Executes a advanced query and returns all objects in the resultset (if any), user can specify the query but it must be correct syntax
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		table_name       
		recordJSON           - advanced query (JSON format)
			e.x. { "field" : { $gt: value } } or { "field" : { $gt: value1, $lt: value2 } }
			for more queries, refer to http://www.mongodb.org/display/DOCS/Advanced+Queries	
        """
        result = None
	dbCollName = str(dbCollName)
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and execute "' + recordJSON + '"'
	    table = self.db[dbCollName]
	    
	    try:
	        recordJSON = dict(json.loads(recordJSON))
	    except:
		print "It's not a valid JSON"
		return result
	    else:
		print "It's a valid JSON"
		coll = self.db['%s' % (dbCollName)]
		result = coll.remove(recordJSON)
		return result
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')

    def AdvanceQueryCount(self, dbCollName, recordJSON):
        """ Executes a advanced query and returns all objects in the resultset (if any), user can specify the query but it must be correct syntax
            Assumes that connection to db has been created with connectToMongoDB.
	    Input:
		dbCollName           - Dababase Collection Name       
		recordJSON           - Advanced JSON Message
			e.x. { "field" : { $gt: value } } or { "field" : { $gt: value1, $lt: value2 } }
			for more queries, refer to http://www.mongodb.org/display/DOCS/Advanced+Queries	
        """
        result = None
        
        if self.conn is not None:
            print 'Executing table: "' + dbCollName + '" and execute "' + recordJSON + '"' 
	    table = self.db[dbCollName]
		
	    try:
		r = dict(json.loads(recordJSON))
	    except:
		print "It's not a valid JSON"
	    else:
	        result = table.find(r).count()
            
	    return result
	
        else:
            asserts.fail('No Mongo db connection! Please make sure to open connection.')
			
			
			
    def SelectTheLastEntry(self, dbCollName, dbFieldName):
		sortOrder = "{\"" + dbFieldName + "\": -1}"
		if self.conn is not None:
			print 'Executing table: "' + dbCollName
			print sortOrder
			table = self.db[dbCollName]
			for record in table.find().sort(sortOrder) :
				result = "%s\n" % (record)
		print result
		return result
			
        